import argparse
import threading
from live_source import LiveSource
from live_source import LiveStreamNotFoundException
import configparser
import subprocess
import logging

# 配置日志记录
logging.basicConfig(
    filename="app.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# 创建控制台处理程序并配置
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
console_handler.setFormatter(console_formatter)

# 将控制台处理程序添加到根记录器
root_logger = logging.getLogger()
root_logger.addHandler(console_handler)


def main():
    try:
        parser = argparse.ArgumentParser(description="抖音直播分发器")
        parser.add_argument("-u", "--url", required=True, help="抖音直播链接")

        args = parser.parse_args()

        live_url = args.url

        try:
            if live_source_url := get_live_source(live_url):
                logging.info(f"直播源解析完毕：{live_source_url}")
                # 开始进行推流
                distribution(live_source_url, get_target_urls())
            else:
                logging.error("解析失败")
                input("按任意键关闭程序...")
        except LiveStreamNotFoundException as e:
            logging.error(f"直播源解析发生错误：{e}")
            input("按任意键关闭程序...")
    except Exception as e:
        logging.error(f"发生错误：{e}")


# 获取源直播流
def get_live_source(live_url):
    logging.info("开始获取源直播流地址...")
    ls = LiveSource()
    return ls.async_url(live_url=live_url)


# 获取分发地址
def get_target_urls():
    config = configparser.ConfigParser()
    config.read("config.ini", encoding="utf-8-sig")
    target_urls_str = config.get("rtmp", "target_urls")
    logging.info(f"获取分发地址：{target_urls_str}")
    return eval(target_urls_str)


# 分发函数
def distribution(live_source_url, target_urls):
    logging.info("============想停止哪个流的直播就关掉对应的ffmpeg窗口============")
    # 启动多个FFmpeg进程来分发到不同的直播流
    for output_stream in target_urls:
        ffmpeg_command = [
            "ffmpeg",
            "-i",
            live_source_url,
            "-c:v",
            "copy",
            "-c:a",
            "copy",
            "-f",
            "flv",
            output_stream,
        ]

        # 启动FFmpeg进程并将其添加到列表中
        process = subprocess.Popen(
            ffmpeg_command,
            universal_newlines=True,
            creationflags=subprocess.CREATE_NEW_CONSOLE,
        )

        threading.Thread(target=task, args=(process, output_stream)).start()


# 任务
def task(process, target_urls):
    logging.info(f"ffmpeg推流任务创建成功，地址：{target_urls}")
    while process.poll() is None:
        pass
    logging.info(f"ffmpeg推流任务停止推送，地址：{target_urls}")


if __name__ == "__main__":
    main()
